import torch
import logging
import editdistance

from shutil import move


class AttrDict(dict):
    """
    Dictionary whose keys can be accessed as attributes.
    """

    def __init__(self, *args, **kwargs):
        super(AttrDict, self).__init__(*args, **kwargs)

    def __getattr__(self, item):
        if item not in self:
            return None
        if type(self[item]) is dict:
            self[item] = AttrDict(self[item])
        return self[item]

    def __setattr__(self, item, value):
        self.__dict__[item] = value


def init_logger(log_file=None):
    log_format = logging.Formatter("[%(asctime)s %(levelname)s] %(message)s")
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)

    console_handler = logging.StreamHandler()
    console_handler.setFormatter(log_format)
    logger.handlers = [console_handler]

    if log_file and log_file != '':
        file_handler = logging.FileHandler(log_file)
        file_handler.setFormatter(log_format)
        logger.addHandler(file_handler)
    return logger


def computer_cer(preds, labels):
    dist = sum(editdistance.eval(label, pred) for label, pred in zip(labels, preds))
    total = sum(len(l) for l in labels)
    return dist, total


def get_saved_folder_name(config):
    return '_'.join([config.data.name, config.training.save_model])


def count_parameters(model):
    n_params = sum([p.nelement() for p in model.parameters()])
    enc = 0
    dec = 0
    for name, param in model.named_parameters():
        if 'encoder' in name:
            enc += param.nelement()
        elif 'decoder' in name:
            dec += param.nelement()
    return n_params, enc, dec


def init_parameters(model, type='xuniform'):
    for p in model.parameters():
        if p.dim() > 1:
            if type == 'xuniform':
                torch.nn.init.xavier_uniform_(p)
            elif type == 'xnoraml':
                torch.nn.init.xavier_normal_(p)
            elif type == 'uniform':
                torch.nn.init.uniform_(p, -0.1, 0.1)

        else:
            pass


def add_space(path, old_path):
    """
    add space between chars

    Arguments
    -------
    path : str
        generate new text to path
    old_path : str
        move old text file to old path
    """
    with open(path, "r") as f_in:
        lines = f_in.readlines()

    move(path, old_path)

    with open(path, "w") as f_out:
        for line in lines:
            data = line.split()
            utt = data[0]
            txt = "".join(data[1:])
            f_out.write(utt + " " + " ".join(list(txt)) + "\n")


def num_gpus(s):
    return len([int(i) for i in s.strip().split(",") if len(i) > 0])


import matplotlib.pyplot as plt
from pylab import *

zhfont1 = matplotlib.font_manager.FontProperties(
    fname="/home1/meichaoyang/workspace/git/kws_ctc_no2/data/SourceHanSansSC-Bold.otf")

"""
usage:
    encoder_output, output_lengths=model.get_post(inputs,inputs_length, True)
    show_ctc_loss(encoder_output[1],targets[1],train_dataset.idx2unit, "exp/png/dnn_0.1.png")
"""


def show_ctc_loss(utt_prob, target, idx2unit, save_path):
    """
    Example
    --------
    >>> encoder_output, output_lengths=model.get_post(inputs,inputs_length, True)
    >>> show_ctc_loss(encoder_output[1],targets[1],train_dataset.idx2unit, "exp/png/dnn_0.1.png")
    """
    plt.plot(utt_prob[:, 0].detach().numpy(), linewidth=0.1)
    for j in target:
        if j != 0:
            plt.plot(utt_prob[:, int(j)].detach().numpy(), linewidth=0.5, linestyle="-.", label=idx2unit[int(j)])

    legend(loc='upper right', prop=zhfont1)
    savefig(save_path, dpi=440)


def show_attention_weights(attention_weights, preds, idx2unit, save_path):
    """
    Example
    --------
    >>> preds = model.recognize(inputs, inputs_length)
    >>> show_attention_weights(model.decoder.attention_weights, preds[-1],
    >>>                dev_dataset.idx2unit, "save_path1.png" )
    """
    for i, ch in enumerate(preds):
        plt.plot(attention_weights[i][0, 0].numpy(), linewidth=0.5, label=idx2unit[ch])

    legend(loc='upper right', prop=zhfont1)
    savefig(save_path, dpi=440)
